# -*- coding: utf-8 -*-

import os

import crypto

try:
    import cPickle as pickle
except ImportError:
    import pickle


se = crypto.safe_encode
sd = lambda x: unicode(x, 'utf-8', errors='replace')


class KeySafe(object):

    def __init__(self, password, path, auto_sync=True):
        self.path = os.path.abspath(path)
        self.auto_sync = auto_sync
        try:
            with open(path, 'rb') as f:
                raw = f.read()
            self._safe = pickle.loads(crypto.decrypt(raw, password))
            del self._safe['_default_']['__rdata']
        except IOError:
            self._safe = {'_default_': {}}

    def __repr__(self):
        return 'KeySafe {0} at {1}'.format(self.path, id(self))

    def __iter__(self):
        for cat in self._safe:
            yield (sd(cat), [sd(x) for x in self._safe[cat].keys()])

    def sync(self, password):
        # We add some random data here to obfuscate the file content
        self._safe['_default_']['__rdata'] = (
            crypto.encrypt(os.urandom(40), password),
            crypto.encrypt(os.urandom(50), password))
        new = '{0}.tmp'.format(self.path)
        with open(new, 'wb') as f:
            f.write(crypto.encrypt(pickle.dumps(self._safe), password))
        if os.path.isfile(self.path):
            os.remove(self.path)
        os.rename(new, self.path)

    def add(self, password, service, username, userpass, category=''):
        if not category:
            category = '_default_'
        else:
            category = se(category)
        if category not in self._safe:
            self._safe[category] = {}
        uname, pw = se(username), se(userpass)
        self._safe[category][se(service)] = (crypto.encrypt(uname, password),
                                             crypto.encrypt(pw, password))
        if self.auto_sync:
            self.sync(password)

    def get(self, password, service, category=''):
        if not category:
            category = '_default_'
        uname, pw = self._safe[se(category)][se(service)]
        return (sd(crypto.decrypt(uname, password)),
                sd(crypto.decrypt(pw, password)))

    def delete(self, service, category=''):
        if not category:
            category = '_default_'
        try:
            del self._safe[se(category)][se(service)]
        except KeyError:
            pass

    def delete_category(self, category):
        cat = se(category)
        if cat == '_default_':
            return
        try:
            del self._safe[cat]
        except KeyError:
            pass
